import tqdm
from typing import List, Tuple

from pylab import *
from matplotlib import rcParams
# 最佳中文配置
rcParams.update({
    'font.family': ['sans-serif'],
    'font.sans-serif': [
        'Noto Sans CJK JP',  # 中文字体（思源黑体）
        'Nimbus Sans',       # 英文字体
        'DejaVu Sans',       # 备用英文字体
        'Arial',             # 备用
        'sans-serif'         # 最终备用
    ],
    'font.size': 10,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'xtick.labelsize': 9,
    'ytick.labelsize': 9,
    'legend.fontsize': 9,
    
    # 数学字体设置（重要！）
    'mathtext.fontset': 'stix',      # 使用STIX数学字体
    'mathtext.default': 'regular',   # 常规样式
})
DPI=600

from .fv import do_fv

from .trace import safe_filter, filtfilt
from .math import norm


def plot_traces_by_subfigures(trace,t, N1=4,N2=6, event_times=[], fig=None) -> Figure:

    if fig is None:
        plt.close('all')
        fig = figure(figsize=(10, 10), dpi=DPI)
    
    N_split = N1*N2
    tn = len(trace)//N_split
    for i in range(N_split):
        axi = subplot(N1,N2,i+1)
        t_i = t[i*tn:(i+1)*tn]
        trace_i = trace[i*tn:(i+1)*tn]
        line, =plot(t_i, trace_i, lw=0.5, color='k', label=f'{t_i[0]:.2f}-{t_i[-1]:.2f}s')
        if len(event_times)>0:
            for t_e in event_times:
                plot([t_e,t_e],[trace_i.min(),trace_i.max()],'r',lw=0.5)
        xlim([t_i.min(),t_i.max()])
        legend(handles=[line], loc='upper right')

    return fig

def plot_events(traces,x,t,fs=1,fe=20, 
                fig=None, ax=None, SCALE=1, DO_FILTER=False,NORM=False, PLOT_WIGGLE=True)-> Tuple[Figure, Axes]:

    if fig is None:
        plt.close('all')
        fig = figure(figsize=(6, 6), dpi=DPI)
    nx,nt = traces.shape
    if NORM:
        traces = norm(traces, ONE_AXIS=True)

    dt = t[1]-t[0]
    if ax is None:
        ax = subplot(111)
    else:
        plt.sca(ax)

    if PLOT_WIGGLE:
        for i in range(nx):
            trace_i = traces[i,:].copy()
            if DO_FILTER:
                trace_i = safe_filter(trace_i, dt, ftype='bandpass', zerophase=True, freqmin=fs, freqmax=fe)
                # trace_i = filtfilt(trace_i, dt, [fs,fe],order=4)
            if NORM:
                trace_i = norm(trace_i)
            trace_i = trace_i*SCALE
            line, =plot(t, trace_i+x[i], lw=0.5, color='k')
        ylim([x.min()-SCALE*traces.max(), x.max()+SCALE*traces.max()])
    else:
        print('not plot wiggle')
        x,d_slices = sort_data_by_distance(x,traces)
        x = np.array(x)
        X,Y = np.meshgrid(t,x)
        pcolormesh(X,Y,d_slices,cmap='RdBu', shading='nearest', rasterized=True)
        # colorbar()

        ylim([x.min(),x.max()])

    xlim([t.min(),t.max()])
    xlabel('t(s)')
    
    ylabel('x')

    return fig, ax

def sort_data_by_distance(distances, DATA=None, NAMES=None):
    ns = len(distances)
    if NAMES is not None:
        assert ns==len(NAMES)
        names=NAMES
    else:
        names=['0']*ns
    if DATA is not None:
        assert ns==DATA.shape[0]
        data=DATA
    else:
        data=np.zeros([ns,1])
    if NAMES is None and DATA is None:
        raise ValueError('At least one para in DATA and NAMES is not None')

    list_data = []
    for i in range(ns):
        trace = data[i,:]
        dist_i = distances[i]
        name_i = names[i]
        list_data.append([dist_i,trace,name_i])
    list_data.sort(key=lambda x:x[0])

    distances_sorted=[]
    data_sorted=np.zeros_like(data)
    names_sorted=[]
    for i in range(ns):
        distances_sorted.append(list_data[i][0])
        data_sorted[i,:]=list_data[i][1]
        names_sorted.append(list_data[i][2])
    if NAMES is None:
        return distances_sorted,data_sorted
    elif DATA is None:
        return distances_sorted,names_sorted
    else:
        return distances_sorted,data_sorted,names_sorted
def plot_raw_data(d_slices,x,t,fs,fe, VLIM=[-500,500], FV_NORM=False, 
                  fig=None, ax1=None, ax2=None, PLOT_WIGGLE=False,SCALE=0):

    if fig is None:
        plt.close('all')
        fig = figure(figsize=(10, 5), dpi=DPI)
    nx,nt  = d_slices.shape
    dt=t[1]-t[0]
    if ax1 is None:
        ax1 = subplot(121)
    else:
        plt.sca(ax1)

    if not PLOT_WIGGLE:
        print('not plot wiggle')
        x,d_slices = sort_data_by_distance(x,d_slices)
        x = np.array(x)
        X,Y = np.meshgrid(t,x)
        pcolormesh(X,Y,d_slices,cmap='RdBu', shading='nearest', rasterized=True)
    else:
        for i in range(nx):
            trace_i = d_slices[i,:].copy()
            trace_i = norm(trace_i)
            trace_i = trace_i*SCALE
            line, =plot(t, trace_i+x[i], lw=0.5, color='k')

    xlim([t.min(),t.max()])
    xlabel('t(s)')
    ylim([x.min()-SCALE, x.max()+SCALE])
    ylabel('x')

    # FV
    if ax2 is None:
        ax2 = subplot(122)
    else:
        plt.sca(ax2)
    vs = np.linspace(VLIM[0],VLIM[1],800)
    FV, f,_ = do_fv(d_slices,vs,x=np.array(x), sx=0,dt=dt, NORM=FV_NORM, PRINT=False, fs=fs, fe=fe)

    FV = np.abs(FV)
    FV = FV/FV.max()

    X,Y = np.meshgrid(f,vs)
    pcolormesh(X,Y,FV.T, cmap='nipy_spectral_r')
    xlabel('Frequency [Hz]')
    ylabel('Velocity [m/s]')
    xlim(fs,fe)

    fig.tight_layout()
    return fig,ax1,ax2

def plot_traces(traces,t, fig=None, NORM=True, x=None):

    if fig is None:
        plt.close('all')
        fig = figure(figsize=[4,4], dpi=600)
    
    ax = subplot(111)
    if NORM:
        traces = norm(traces, ONE_AXIS=True)
    if x is None:
        x = np.arange(traces.shape[0])
    X,Y = np.meshgrid(t, x)
    pcolormesh(X,Y, traces, cmap='bwr', shading='nearest', rasterized=True )
        
    xlim([t.min(),t.max()])
    xlabel('t(s)')
    ylim([x.min(), x.max()])
    ylabel('trace number')
    
    return fig, ax
    
def plot_scatters(x,y, fig=None, XL='x', YL='y',mcolor='tab:blue'):
    if fig is None:
        plt.close('all')
        fig = figure(figsize=[8,2.5], dpi=600)
    
    ax = subplot(111)
    scatter(x,y, s=10, c=mcolor, marker='^', edgecolors='none')

    xlim([x.min(),x.max()])
    ylim([y.min(), y.max()])
    xlabel(XL)
    ylabel(YL)

    return fig, ax

def test():
    def get_ap_level(par_name,ts,events_time,folder='./data/GLDAS'):     
        a = np.loadtxt(f'{folder}/{par_name}.txt')
        ap = a[:,1]
        time_ap_sec=[]

        for i in a[:,0]:
            time_i = datetime.datetime.strptime(str(int(i)), '%Y%m%d%H')
            delta =  time_i-ts
            time_ap_sec.append(delta.seconds +delta.days*3600*24)
        
        
        f_ap = interpolate.interp1d(time_ap_sec,ap )
        ap_level_true = f_ap(events_time)
        
        return ap_level_true,ap_level_true

    ts = datetime.datetime.strptime('2023.03.10', '%Y.%m.%d')  # start date and time
    ap_level,_=get_ap_level('GTide',ts,te_after_stack)
    ap_level = -ap_level/ap_level.max()*20